#!/usr/bin/env python3
"""
Semantic Point Cloud Segmentation Utility

This module provides semPcd function for segmenting segpointcloud in unifiedData
into semantic parts according to allowed_seg_colors.json's indices2obj for the current embodiment.
"""

import json
import numpy as np
from pathlib import Path
from typing import Dict, List
from termcolor import colored
from loadJson import _load_seg_json, _load_emb_name_map
from time_it import timeit
@timeit()
def semPcd(unifiedData: dict, embodiment: str, surpress_print: bool = False, return_colors: bool = False, seg_source="traverse") -> Dict[str, List[int]]:
    """
    Segment segpointcloud in unifiedData into semantic parts and return a dictionary
    mapping semantic part names to lists of point indices.
    
    Args:
        unifiedData: Dictionary containing segpointcloud data
        embodiment: Single character embodiment code (e.g., 'f' for franka-panda)
        surpress_print: Whether to suppress print statements
        return_colors: Whether to also return top_indices2obj for target colors
        seg_source: Whether this is called from the pipeline (affects JSON file selection)
        
    Returns:
        Dictionary with semantic part names as keys and lists of point indices as values.
        If return_colors=True, also returns top_indices2obj in a tuple (semantic_indices, top_indices2obj).
    """
    
    seg_colors = _load_seg_json(seg_source)
    emb_map = _load_emb_name_map()
    robot_name = emb_map.get(embodiment, embodiment)
    robot_indices2obj = seg_colors.get(robot_name, {}).get('indices2obj', {})
    top_indices2obj = seg_colors.get('indices2obj', {})
    assert robot_indices2obj, f"semPcd: No indices2obj found for robot {robot_name}"
    assert 'segpointcloud' in unifiedData and unifiedData['segpointcloud'] is not None, "semPcd: No segpointcloud to segment."
    segpcd = unifiedData['segpointcloud']
    assert segpcd.shape[0] > 0, "semPcd: segpointcloud is empty."
    semantic_colors = {}
    for part, robot_val in robot_indices2obj.items():
        indices = robot_val if isinstance(robot_val, list) else [robot_val]
        part_colors = []
        for idx in indices:
            found = False
            for group in ['env', 'gripper']:
                groupinfo = seg_colors.get(robot_name, {}).get(group, {})
                indices_list = groupinfo.get('indices', [])
                colors_list = groupinfo.get('colors', [])
                if idx in indices_list:
                    idx_in_group = indices_list.index(idx)
                    if idx_in_group < len(colors_list):
                        src_color = tuple(colors_list[idx_in_group])
                        part_colors.append(src_color)
                        found = True
                        break
            if not found and not surpress_print:
                print(f"semPcd: Could not find color for index {idx} in {robot_name}")
        if part_colors:
            semantic_colors[part] = part_colors
    semantic_indices = {}
    total_points = 0
    for semantic_part, colors in semantic_colors.items():
        part_indices = []
        for color in colors:
            mask = np.all(segpcd[:, 3:6] == color, axis=1)
            color_indices = np.where(mask)[0].tolist()
            part_indices.extend(color_indices)
        if part_indices:
            semantic_indices[semantic_part] = part_indices
            total_points += len(part_indices)
            if not surpress_print:
                print(f"    🔍 SEM | {semantic_part}: {len(part_indices)} points")
    if not surpress_print:
        print(f"    📊 SEM | Total {len(semantic_indices)} semantic parts, {total_points} points segmented")
    if return_colors:
        return semantic_indices, top_indices2obj
    else:
        return semantic_indices 